"""Bounding box utility functions."""
import numpy as np
from typing import List, Tuple, Union


def compute_iou(box1: Union[List, np.ndarray], box2: Union[List, np.ndarray]) -> float:
    """
    Compute IoU between two boxes.

    Args:
        box1: [x1, y1, x2, y2] format
        box2: [x1, y1, x2, y2] format

    Returns:
        IoU value between 0 and 1
    """
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])

    inter = max(0, x2 - x1) * max(0, y2 - y1)
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union = area1 + area2 - inter

    return inter / (union + 1e-6)


def compute_iou_batch(boxes1: np.ndarray, boxes2: np.ndarray) -> np.ndarray:
    """
    Compute IoU between two sets of boxes.

    Args:
        boxes1: Nx4 array of boxes in [x1, y1, x2, y2] format
        boxes2: Mx4 array of boxes in [x1, y1, x2, y2] format

    Returns:
        NxM array of IoU values
    """
    # Expand dimensions for broadcasting
    boxes1 = np.expand_dims(boxes1, 1)  # Nx1x4
    boxes2 = np.expand_dims(boxes2, 0)  # 1xMx4

    # Compute intersection coordinates
    x1 = np.maximum(boxes1[..., 0], boxes2[..., 0])
    y1 = np.maximum(boxes1[..., 1], boxes2[..., 1])
    x2 = np.minimum(boxes1[..., 2], boxes2[..., 2])
    y2 = np.minimum(boxes1[..., 3], boxes2[..., 3])

    # Compute intersection area
    inter = np.maximum(0, x2 - x1) * np.maximum(0, y2 - y1)

    # Compute areas
    area1 = (boxes1[..., 2] - boxes1[..., 0]) * (boxes1[..., 3] - boxes1[..., 1])
    area2 = (boxes2[..., 2] - boxes2[..., 0]) * (boxes2[..., 3] - boxes2[..., 1])

    # Compute union and IoU
    union = area1 + area2 - inter
    iou = inter / (union + 1e-6)

    return np.squeeze(iou)


def xywh_to_xyxy(boxes: np.ndarray, image_width: int, image_height: int) -> np.ndarray:
    """
    Convert YOLO format [cx, cy, w, h] (normalized) to [x1, y1, x2, y2] (absolute).

    Args:
        boxes: Nx4 array in [cx_norm, cy_norm, w_norm, h_norm] format
        image_width: Image width in pixels
        image_height: Image height in pixels

    Returns:
        Nx4 array in [x1, y1, x2, y2] format
    """
    boxes = np.array(boxes)
    if len(boxes.shape) == 1:
        boxes = boxes.reshape(1, -1)

    result = np.zeros_like(boxes)
    result[:, 0] = (boxes[:, 0] - boxes[:, 2] / 2) * image_width   # x1
    result[:, 1] = (boxes[:, 1] - boxes[:, 3] / 2) * image_height  # y1
    result[:, 2] = (boxes[:, 0] + boxes[:, 2] / 2) * image_width   # x2
    result[:, 3] = (boxes[:, 1] + boxes[:, 3] / 2) * image_height  # y2

    return result


def xyxy_to_xywh(boxes: np.ndarray, image_width: int, image_height: int) -> np.ndarray:
    """
    Convert [x1, y1, x2, y2] (absolute) to YOLO format [cx, cy, w, h] (normalized).

    Args:
        boxes: Nx4 array in [x1, y1, x2, y2] format
        image_width: Image width in pixels
        image_height: Image height in pixels

    Returns:
        Nx4 array in [cx_norm, cy_norm, w_norm, h_norm] format
    """
    boxes = np.array(boxes)
    if len(boxes.shape) == 1:
        boxes = boxes.reshape(1, -1)

    result = np.zeros_like(boxes, dtype=np.float32)
    result[:, 0] = ((boxes[:, 0] + boxes[:, 2]) / 2) / image_width   # cx_norm
    result[:, 1] = ((boxes[:, 1] + boxes[:, 3]) / 2) / image_height  # cy_norm
    result[:, 2] = (boxes[:, 2] - boxes[:, 0]) / image_width         # w_norm
    result[:, 3] = (boxes[:, 3] - boxes[:, 1]) / image_height        # h_norm

    return result


def clip_boxes(boxes: np.ndarray, image_width: int, image_height: int) -> np.ndarray:
    """
    Clip boxes to image boundaries.

    Args:
        boxes: Nx4 array in [x1, y1, x2, y2] format
        image_width: Image width
        image_height: Image height

    Returns:
        Clipped boxes
    """
    boxes = np.array(boxes)
    boxes[:, 0] = np.clip(boxes[:, 0], 0, image_width)
    boxes[:, 1] = np.clip(boxes[:, 1], 0, image_height)
    boxes[:, 2] = np.clip(boxes[:, 2], 0, image_width)
    boxes[:, 3] = np.clip(boxes[:, 3], 0, image_height)

    return boxes


def validate_boxes(boxes: Union[List, np.ndarray]) -> bool:
    """
    Validate box coordinates.

    Args:
        boxes: Single box [x1, y1, x2, y2] or Nx4 array

    Returns:
        True if valid, False otherwise
    """
    boxes = np.array(boxes)
    if len(boxes.shape) == 1:
        boxes = boxes.reshape(1, -1)

    # Check dimensions
    if boxes.shape[1] != 4:
        return False

    # Check x1 < x2 and y1 < y2
    if np.any(boxes[:, 0] >= boxes[:, 2]) or np.any(boxes[:, 1] >= boxes[:, 3]):
        return False

    # Check non-negative coordinates
    if np.any(boxes < 0):
        return False

    return True


def compute_box_area(boxes: np.ndarray) -> np.ndarray:
    """
    Compute area of boxes.

    Args:
        boxes: Nx4 array in [x1, y1, x2, y2] format

    Returns:
        N array of areas
    """
    boxes = np.array(boxes)
    if len(boxes.shape) == 1:
        boxes = boxes.reshape(1, -1)

    return (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
